#include "cvxEdgeShadingRecovery.h"
//#include "cvxLeastSquare.h"
#include <map>

using std::map;
using cvx::Rgb8UImage;
using cvx::Bw8UImage;
using cvx::Bw32FImage;
using cvx::Rgb32FImage;
using cvx::Yuv8UImage;
using cvx::Bw32SImage;

CvxEdgeShadingRecovery::CvxEdgeShadingRecovery()
{
}

CvxEdgeShadingRecovery::~CvxEdgeShadingRecovery()
{
}

void CvxEdgeShadingRecovery::setImage(const Mat &orgImage)
{
	assert(!orgImage.empty());

	int h = orgImage.rows;
	int w = orgImage.cols;

	//normalized image
	m_orgImage = orgImage.clone();
	m_orgImage.convertTo(m_orgImage, CV_32FC3, 1.0/255.0, 0.0002);

	Mat yuvImage = orgImage.clone();

	cv::cvtColor(m_orgImage, yuvImage, CV_BGR2YUV);

	// convert to YUV and 32 bit float
	m_yImage = cv::Mat::zeros(h, w, CV_32FC1);

	Yuv8UImage yuv_img(&yuvImage);
	Bw32FImage y_img(&m_yImage);
	
	for(int y = 0; y<h; y++)
	{
		for(int x = 0; x<w; x++)
		{
			y_img[y][x] = yuv_img[y][x].y/255.0 + 0.0002;
		}
	}

	m_shadingMap = cv::Mat::zeros(h, w, CV_32FC1);
	m_albedoMap  = cv::Mat::zeros(h, w, CV_32FC3);
	m_weightMap  = cv::Mat::zeros(h, w, CV_32FC(9));

}

void CvxEdgeShadingRecovery::calculateWeight(bool isShow)
{
	assert(!m_weightMap.empty());
	assert(!m_orgImage.empty());
	assert(!m_yImage.empty());

	int h = m_orgImage.rows;
	int w = m_orgImage.cols;

	Bw32FImage  y_img(&m_yImage);

	//ignore boundary pixels
	for(int y = 1; y<h-1; y++)
	{
		for(int x = 1; x<w-1; x++)
		{
			//calculate variance
			Scalar mean, stdDev;
			cv::Rect r1 = cv::Rect(x-1, y-1, 3, 3);
			Mat m1 = m_yImage(r1);						
			cv::meanStdDev(m1, mean, stdDev);

			if(stdDev[0] < 0.0002)
			{
				stdDev[0] = 0.0002;
			}

		    //	int test = 1;
			float *wt = (float *)(m_weightMap.ptr(y, x));
			float mid_val = y_img[y][x];
			float sum_wt = 0.0;
			for(int i = -1; i <= 1; i++)
			{
				for(int j = -1; j <= 1; j++)
				{
					int idx = (i+1) * 3 + (j+1);
					if(i == 0 && j == 0)
					{
						wt[idx] = 0.0;
						continue;
					}
					float cur_val = y_img[y+i][x+j];
					wt[idx] = exp(-1.0 * (cur_val - mid_val) * (cur_val - mid_val)/(2.0 * stdDev[0] * stdDev[0]));
					sum_wt += wt[idx];
				}
			}

			//normalize weight
			for(int i = 0; i<9; i++)
			{
				wt[i] = wt[i]/sum_wt;
			}
		}
	}
}

void CvxEdgeShadingRecovery::testWeightWithGroundTruth(Mat &albedoGroundTruth)
{
	int h = m_orgImage.rows;
	int w = m_orgImage.cols;

	Rgb8UImage albedo_gd(&albedoGroundTruth);
	for(int y =0 ; y<h; y++)
	{
		for(int x = 0; x<w; x++)
		{
			float *wt = (float *)(m_weightMap.ptr(y, x));
			float accu_val = 0.0;
			for(int i = -1; i<=1; i++)
			{
				for(int j = -1; j<=1; j++)
				{
					if(i == 0 && j == 0)
					{
						continue;
					}

					int cy = y+i;
					int cx = x+j;						
					int wtIdx = (i+1)*3 + (j+1);
					accu_val += wt[wtIdx] * albedo_gd[cy][cx].r;
					printf("%f ", wt[wtIdx]);
				}
			}
			printf("residence is %f\n", abs(albedo_gd[y][x].r - accu_val));
		}
	}

}

void CvxEdgeShadingRecovery::iterativeSolve(int iterNum, float lambda)
{
	assert(!m_orgImage.empty());
	assert(!m_shadingMap.empty());	
	assert(!m_albedoMap.empty());

	int h = m_orgImage.rows;
	int w = m_orgImage.cols;

	//initialize shading to 1.0
	Bw32FImage  rev_sha_img(&m_shadingMap);
	Rgb32FImage org_img(&m_orgImage);
	Rgb32FImage alb_img(&m_albedoMap);
	for(int y = 0; y<h; y++)
	{
		for(int x = 0; x<w; x++)
		{
			//0.3 0.59 0.11
			rev_sha_img[y][x] = 2.0;
			alb_img[y][x].r = 0.5;
			alb_img[y][x].g = 0.5;
			alb_img[y][x].b = 0.5;
		}
	}


	for(int k = 0; k<iterNum; k++)
	{
		//update albedo
		//sumR = W(len,:)*R(n_idx(len,:),:) + sI(len,:);
		//R(len,:) = (1-rho)*R(len,:) + 0.5*rho*sumR;
		printf("k = %d\n", k);

		Mat secondAlbedo = m_albedoMap.clone();
		Rgb32FImage sec_alb(&secondAlbedo);
		for(int y = 1; y<h-1; y++)
		{
			for(int x = 1; x<w-1; x++)
			{
				float sumR[3] = {0.0, 0.0, 0.0};
				float *wt = (float *)(m_weightMap.ptr(y, x));
				//from neighbors
				for(int i = -1; i<=1; i++)
				{
					for(int j = -1; j<=1; j++)
					{
						int w_idx = (i+1) * 3 + j+1;
						int cy = y + i;
						int cx = x + j;
						sumR[0] += wt[w_idx] * alb_img[cy][cx].r;
						sumR[1] += wt[w_idx] * alb_img[cy][cx].g;
						sumR[2] += wt[w_idx] * alb_img[cy][cx].b;
					}
				}

				float r = org_img[y][x].r;
				float g = org_img[y][x].g;
				float b = org_img[y][x].b;
				float s_rv = rev_sha_img[y][x];
				
				//from I
				sumR[0] += org_img[y][x].r * rev_sha_img[y][x];
				sumR[1] += org_img[y][x].g * rev_sha_img[y][x];
				sumR[2] += org_img[y][x].b * rev_sha_img[y][x];

				//0.5 for 2 (1 for neighbor, 1 for I)
				
				sec_alb[y][x].r = (1.0 - lambda) * alb_img[y][x].r + 0.5 * lambda * sumR[0];
				sec_alb[y][x].g = (1.0 - lambda) * alb_img[y][x].g + 0.5 * lambda * sumR[1];
				sec_alb[y][x].b = (1.0 - lambda) * alb_img[y][x].b + 0.5 * lambda * sumR[2];
			
				float epsilon = 0.001;
				if(sec_alb[y][x].r < epsilon)
				{
					sec_alb[y][x].r = epsilon;
				}
				if(sec_alb[y][x].r > 1.0)
				{
					sec_alb[y][x].r = 1.0;
				}

				if(sec_alb[y][x].g < epsilon)
				{
					sec_alb[y][x].g = epsilon;
				}
				if(sec_alb[y][x].g > 1.0)
				{
					sec_alb[y][x].g = 1.0;
				}

				if(sec_alb[y][x].b < epsilon)
				{
					sec_alb[y][x].b = epsilon;
				}
				if(sec_alb[y][x].b > 1.0)
				{
					sec_alb[y][x].b = 1.0;
				}
				
			}
		}

		//update inv_sha
		//	 inv_S = (1-rho)*inv_S + rho*(sum(rgbIm.*R,2)./rgbIm2);
	//     inv_S = max(1, inv_S);
		for(int y = 1; y<h-1; y++)
		{
			for(int x = 1; x<w-1; x++)
			{
				float r = sec_alb[y][x].r/org_img[y][x].r;
				float g = sec_alb[y][x].g/org_img[y][x].g;
				float b = sec_alb[y][x].b/org_img[y][x].b;

				float revS = sqrt(r*r + g*g + b*b);
				rev_sha_img[y][x] = (1.0 - lambda) * rev_sha_img[y][x] + lambda * revS;
				if(rev_sha_img[y][x] < 1.0)
				{
					rev_sha_img[y][x] = 1.0;
				}
			}
		}

		secondAlbedo.copyTo(m_albedoMap);


		{
			Mat showAlbedo = m_albedoMap.clone();
			showAlbedo.convertTo(showAlbedo, CV_8UC3, 255.0);

			cv::imshow("albedo", showAlbedo); 
			Mat showShading = cv::Mat::zeros(h, w, CV_8UC1);
			Bw8UImage show_sha(&showShading);
			for(int y = 1; y<h-1; y++)
			{
				for(int x = 1; x<w-1; x++)
				{
					show_sha[y][x] = 1.0/rev_sha_img[y][x] * 255.0;
				}
			}
			cv::imshow("shading", showShading);
			cv::waitKey(0);
		}
	}

	//show result
	{
		Mat showAlbedo = m_albedoMap.clone();
		showAlbedo.convertTo(showAlbedo, CV_8UC3, 255.0);

		cv::imshow("albedo", showAlbedo); 
		Mat showShading = cv::Mat::zeros(h, w, CV_8UC1);
		Bw8UImage show_sha(&showShading);
		for(int y = 1; y<h-1; y++)
		{
			for(int x = 1; x<w-1; x++)
			{
				show_sha[y][x] = 1.0/rev_sha_img[y][x] * 255.0;
			}
		}
		cv::imshow("shading", showShading);
	}
}

